import commands
import numpy as np
import StringIO
from time import time

class PolymorphismTable:
	def __init__( self, positionVector, dataMatrix ):
		self.positions = positionVector
		self.table = dataMatrix

	def __str__(self):
		return '\nPolymorphism table: \n%s' % str(self.table)

class PolymorphismTables:
	def __init__(self):
		self.polyTables = []

	def __str__(self):
		s = ''
		for p in self.polyTables:
			s += str(p)
		return s

	def __iter__(self):
		self.itr = iter(self.polyTables)
		return self

	def next(self):
		return self.itr.next()

	def runMS( self, parameters, outFile='/scratch/damdavid/ms.out' ):
		newCommand = '~/bin/ms' + parameters + ' -seeds 1 1 1'
		print '\nRunning\n', newCommand
		msOut = commands.getoutput( newCommand )
		
#		outFile += "%f" % ( time() )
#		
#		f = open( outFile, 'w')
#		f.write( msOut )
#		f.close()
#		
#		commands.getoutput('gzip ' + outFile)
		
		return msOut

	def readMS( self, filename, alreadyReady=None ):
		f = None
		
		if alreadyReady == None:
			f = open( filename, 'r')
		else:
			f = StringIO.StringIO(alreadyReady)
	
		header = f.readline().split()
		numChromosomes = int(header[1])
		numSimulations = int(header[2])
		
		f.readline() #ignore
		f.readline()
		f.readline()
		
		self.polyTables = []
		
		for simItr in range(numSimulations):
			
			if ((simItr + 1) % 1000) == 0:
				print "... Processed 1000 simulations: now at " + str(simItr + 1)
			
			segLine = f.readline().split()
			numSegSites = int(segLine[1])
			
			if numSegSites > 0:
				posLine = f.readline()
				a = []
				for chrItr in range(numChromosomes):
					chr = str(f.readline().rstrip())
					polys = []
					for siteItr in range( len(chr) ):
						polys.append( int(chr[siteItr]) )
	
					a.append( np.array( polys ))
	
				data = np.array(a)
				poly = PolymorphismTable(posLine, data)
				
				self.polyTables.append( poly )
			
			else:
				f.readline()
			
			f.readline()
			f.readline()
		
		return self.polyTables

